[Python] pytestのコマンド引数で環境変数を切り替えてみる

どうも！大阪オフィスの西村祐二です。

テストコードを実行するときに環境変数を切り替えたい場面があるかと思います。

例えば、DynamoDBのエンドポイントをLocalstackで立てたエンドポイントに向けたいなど

pytestでそれを行う方法を備忘録かねてブログとしておきます。

やりたいこと

$ pytest test_env.py --local
localstackなどで使う環境変数を設定

$ pytest test_env.py --env=stg
ENVという環境変数を「stg」に設定

$ pytet test_env.py --local --env=stg
上記の2つの環境変数設定

$ pytest test_env.py
デフォルトの環境変数を設定

ENVをわけている理由としては「itg-UserTable」のように(環境名)-(リソース名)というように環境ごとにプレフィックスをつけることが多いためです。

やってみる

環境

• Python 3.7.2

• pytest-4.2.1

環境ごとの環境変数を記載したファイルを作成

デフォルト時に読み込む環境変数を記載

{
    "DYNAMODB_ENDPOINT_URL": "https://dynamodb.ap-northeast-1.amazonaws.com/"
}

--localという引数があるときに設定する環境変数を記載

{
    "DYNAMODB_ENDPOINT_URL": "http://localhost:4569"
}

pytest

すべてのテストで設定したいので、conftestに記載していきます。

import json
import os
import pytest


def pytest_addoption(parser):
    """Add pytest command options."""
    parser.addoption(
        "--local",
        action="store_true",
        default=False,
        help="set local endpoint"
    )
    parser.addoption(
        '--env',
        action='store',
        default='itg',
        help='set env'
    )


@pytest.fixture(scope='session', autouse=True)
def cmdopt(request):
    """Get options value."""

    # ENV
    os.environ['ENV'] = request.config.getoption("--env")

    # load file
    if(request.config.getoption("--local")):
        set_environment_variable('./variable.local.json')
    else:
        set_environment_variable('./variable.json')


def set_environment_variable(json_filepath):
    """Jsonを読み込み環境変数にセットする."""
    with open(f"./{json_filepath}") as filedata:
        for key, value in json.load(filedata).items():
            os.environ[key] = value

簡単に解説すると

6行目：pytest_addoptionでpytestでコマンド引数を設定します。--localと--envが使えるように設定しています。action="store_true"とすることで、引数があるときtrueを返してくれます。また、--envは引数にないときは「itg」をデフォルトで設定するようにしています。

23行目からのcmdopt関数では、request.config.getoption("")でコマンド引数の値を取得することができるので、取得した値から制御しています。--localがあれば、ローカル用の環境変数を読み込み、なれけば、通常用の環境変数を読み込みます。

36行目からのset_environment_variable関数はファイルから環境変数にセットする処理をしています。

テスト

きちんと環境変数が設定されているかのテストを書いてみます。

▼--localを引数に設定してテストコード実行してみます。

import pytest
import os
def test_env_check():
    """Check environment variable."""
    ENV = os.getenv('ENV')
    DYNAMODB_ENDPOINT_URL = os.getenv('DYNAMODB_ENDPOINT_URL')
    assert 'itg' == ENV
    assert 'http://localhost:4569' == DYNAMODB_ENDPOINT_URL

問題なくテストが通りましたね。

$ pytest test_env.py --local

======================================== test session starts ========================================
platform darwin -- Python 3.7.2, pytest-4.2.1, py-1.7.0, pluggy-0.8.1
rootdir: /Users/nishimura.yuji/study/python/pytest, inifile:
collected 1 item                                                                                    

test_env.py .                                                                                 [100%]

▼次はhttp://localhost:4569を"https://dynamodb.ap-northeast-1.amazonaws.com/に変えて--localの引数はなしでテストコード実行してみます。

import pytest
import os
def test_env_check():
    """Check environment variable."""
    ENV = os.getenv('ENV')
    DYNAMODB_ENDPOINT_URL = os.getenv('DYNAMODB_ENDPOINT_URL')
    assert 'itg' == ENV
    assert 'https://dynamodb.ap-northeast-1.amazonaws.com/' == DYNAMODB_ENDPOINT_URL

こちらも問題なくテストが通りましたね。うまく環境変数が切り替わっています。

$ pytest test_env.py
============================== test session starts ==============================
platform darwin -- Python 3.7.2, pytest-4.2.1, py-1.7.0, pluggy-0.8.1
rootdir: /Users/nishimura.yuji/study/python/pytest, inifile:
collected 1 item                                                                

test_env.py .                                                             [100%]

▼次は、--env=stgの引数をつけて、テストコードをitgからstgに変えて実行してみます。

import pytest
import os
def test_env_check():
    """Check environment variable."""
    ENV = os.getenv('ENV')
    DYNAMODB_ENDPOINT_URL = os.getenv('DYNAMODB_ENDPOINT_URL')
    assert 'stg' == ENV
    assert 'https://dynamodb.ap-northeast-1.amazonaws.com/' == DYNAMODB_ENDPOINT_URL

こちらも問題なくテストが通りましたね。うまく環境変数が切り替わっています。

$ pytest test_env.py --env=stg
============================== test session starts ==============================
platform darwin -- Python 3.7.2, pytest-4.2.1, py-1.7.0, pluggy-0.8.1
rootdir: /Users/nishimura.yuji/study/python/pytest, inifile:
collected 1 item                                                                

test_env.py .                                                             [100%]

▼さいごに、引数に--localと--env=stgをつけて、テストコードはhttp://localhost:4569に変更して実行してみます。

import pytest
import os
def test_env_check():
    """Check environment variable."""
    ENV = os.getenv('ENV')
    DYNAMODB_ENDPOINT_URL = os.getenv('DYNAMODB_ENDPOINT_URL')
    assert 'stg' == ENV
    assert 'http://localhost:4569' == DYNAMODB_ENDPOINT_URL

こちらも問題なくテストが通りましたね。うまく環境変数が切り替わっています。

$ pytest test_env.py --env=stg --local
============================== test session starts ==============================
platform darwin -- Python 3.7.2, pytest-4.2.1, py-1.7.0, pluggy-0.8.1
rootdir: /Users/nishimura.yuji/study/python/pytest, inifile:
collected 1 item                                                                

test_env.py .                                                             [100%]

さいごに

いかがだったでしょうか。

pytestのコマンド引数で環境変数を切り替えてみました。

もっといいやり方がある、こうしたほうがいいなどありましたらコメントいただけると幸いです。